package org.example.controller;


import brave.Span;
import brave.Tracer;
import brave.Tracing;
import brave.propagation.CurrentTraceContext;
import brave.propagation.Propagation;
import brave.propagation.TraceContext;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

@RestController
public class ConsumerController {
    @Autowired
    Tracing tracing;

    Propagation.RemoteGetter<HttpServletRequest> GETTER = new Propagation.RemoteGetter<HttpServletRequest>() {
        @Override
        public Span.Kind spanKind() {
            return Span.Kind.SERVER;
        }

        @Override
        public String get(HttpServletRequest request, String name) {
            return request.getHeader(name);
        }
    };


    @GetMapping("consume1")
    public String consume1(HttpServletRequest request, HttpServletResponse response) {
        CurrentTraceContext currentTraceContext = tracing.currentTraceContext();
        String result = null;
        TraceContext context = (TraceContext) request.getAttribute(TraceContext.class.getName());
        if (context != null) {
            CurrentTraceContext.Scope scope = currentTraceContext.maybeScope(context);
            try {
                result = process();
            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                scope.close();
            }
            return result;
        }

        TraceContext.Extractor<HttpServletRequest> extractor = tracing.propagation().extractor(GETTER);
        Tracer tracer = tracing.tracer();
        Span span = tracer.nextSpan(extractor.extract(request)).name("consume1");
        span.start();
        CurrentTraceContext.Scope scope = currentTraceContext.newScope(span.context());
        try {
            result = process();
        } catch (Exception e) {
            span.error(e);
            e.printStackTrace();
        } finally {
            span.kind(Span.Kind.SERVER);
            span.finish();
            scope.close();
        }
        return result;
    }

    @GetMapping("consume2")
    public String consume2() {
        return process();
    }

    Propagation.RemoteSetter<HttpUriRequest> SETTER = new Propagation.RemoteSetter<HttpUriRequest>() {
        @Override
        public Span.Kind spanKind() {
            return Span.Kind.CLIENT;
        }

        @Override
        public void put(HttpUriRequest request, String name, String value) {
            request.setHeader(name, value);
        }
    };

    /* client */
    private String process() {
        TraceContext.Injector<HttpUriRequest> injector = tracing.propagation().injector(SETTER);
        CurrentTraceContext currentTraceContext = tracing.currentTraceContext();
        Tracer tracer = tracing.tracer();
        Span span = tracer.nextSpan().name("consumer-process").start(); // start a span
        try (CurrentTraceContext.Scope ws = currentTraceContext.newScope(span.context())) {
            CloseableHttpClient httpClient = HttpClients.custom().build();
            HttpGet httpGet = new HttpGet("http://localhost:8081/provide");
            injector.inject(span.context(), httpGet); // inject span context into HTTP!!!
            span.kind(Span.Kind.CLIENT);
            CloseableHttpResponse response = httpClient.execute(httpGet); // RCP
            System.out.println("Response: " + response.getStatusLine());
            System.out.println(EntityUtils.toString(response.getEntity()));
            response.close();
            httpClient.close();
        } catch (Exception e) {
            span.error(e); // deal with error
            e.printStackTrace();
        } finally {
            span.finish(); // finish a span. // note the scope is independent of the span. Always finish a span.
        }
        return "consumer server processed";
    }
}
